// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <turbo/utility/status.h>
#include <nebula/types/type_fwd.h>
#include <turbo/log/logging.h>
#include <turbo/base/macros.h>

namespace nebula {
namespace util {

template <typename T>
void CheckAlignment(const void* ptr) {
  DKCHECK(reinterpret_cast<uint64_t>(ptr) % sizeof(T) == 0);
}

/// Storage used to allocate temporary vectors of a batch size.
/// Temporary vectors should resemble allocating temporary variables on the stack
/// but in the context of vectorized processing where we need to store a vector of
/// temporaries instead of a single value.
class TURBO_EXPORT TempVectorStack {
  template <typename>
  friend class TempVectorHolder;

 public:
  TempVectorStack() = default;
  ~TempVectorStack();

  TURBO_DISALLOW_COPY_AND_ASSIGN(TempVectorStack);

  TURBO_DEFAULT_MOVE_AND_ASSIGN(TempVectorStack);

  turbo::Status init(MemoryPool* pool, int64_t size);

  int64_t AllocatedSize() const { return top_; }

 private:
  static int64_t EstimatedAllocationSize(int64_t size) {
    return PaddedAllocationSize(size) + /*two guards*/ 2 * sizeof(uint64_t);
  }

  static int64_t PaddedAllocationSize(int64_t num_bytes);

  void alloc(uint32_t num_bytes, uint8_t** data, int* id);
  void release(int id, uint32_t num_bytes);
  static constexpr uint64_t kGuard1 = 0x3141592653589793ULL;
  static constexpr uint64_t kGuard2 = 0x0577215664901532ULL;
  static constexpr int64_t kPadding = 64;
  int num_vectors_;
  int64_t top_;
  std::unique_ptr<Buffer> buffer_;
  int64_t buffer_size_;

  friend class TempVectorStackTest;
};

template <typename T>
class TempVectorHolder {
  friend class TempVectorStack;

 public:
  ~TempVectorHolder() { stack_->release(id_, num_elements_ * sizeof(T)); }
  T* mutable_data() { return reinterpret_cast<T*>(data_); }
  TempVectorHolder(TempVectorStack* stack, uint32_t num_elements) {
    stack_ = stack;
    num_elements_ = num_elements;
    stack_->alloc(num_elements * sizeof(T), &data_, &id_);
  }

 private:
  TempVectorStack* stack_;
  uint8_t* data_;
  int id_;
  uint32_t num_elements_;
};

}  // namespace util
}  // namespace nebula
